package cn.yuanqiao.archive.utils.impl;

import cn.yuanqiao.archive.utils.FilePageCounter;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.pdmodel.PDPage;
import org.springframework.stereotype.Service;

import java.io.*;
import java.util.ArrayList;
import java.util.List;

/**
 * @ClassName: PdfFilePageCounter
 * @Description: 类的描述
 * @author: lihuahua
 * @date: 2024/4/11 17:20
 * @Blog: huahuacoco.top
 */

@Service("PdfFilePageCounter")
public class PdfFilePageCounter implements FilePageCounter {

    @Override
    public int getPageCount(String filePath) throws Exception {
        try (PDDocument document = PDDocument.load(new File(filePath))) {
            return document.getNumberOfPages();
        }
    }

    @Override
    public InputStream extractPages(String url, List<Integer> pageNumbers) throws Exception {
        try (PDDocument document = PDDocument.load(new File(url))) {
            List<PDPage> pagesToExtract = new ArrayList<>();
            for (int pageNum : pageNumbers) {
                if (pageNum <= 0 || pageNum > document.getNumberOfPages()) {
                    throw new IllegalArgumentException("Invalid page number: " + pageNum);
                }
                pagesToExtract.add(document.getDocumentCatalog().getPages().get(pageNum - 1));
            }

            // 创建用于存储裁剪后页面的新文档
            PDDocument extractedDocument = new PDDocument();
            for (PDPage page : pagesToExtract) {
                extractedDocument.addPage(page);
            }

            // 将裁剪后的PDF写入内存输出流
            try (OutputStream outputStream = new ByteArrayOutputStream()) {
                extractedDocument.save(outputStream);
                extractedDocument.close();

                return new ByteArrayInputStream(((ByteArrayOutputStream) outputStream).toByteArray());
            }
        }
    }
}
